from dataclasses import dataclass
from ipaddress import IPv4Address
from logging import getLogger
from os.path import dirname
from typing import Iterator, Tuple

from pysnmp.hlapi import (
    CommunityData,
    ContextData,
    ObjectIdentity,
    ObjectType,
    SnmpEngine,
    UdpTransportTarget,
    getCmd,
    setCmd,
)
from pysnmp.proto.errind import ErrorIndication, requestTimedOut
from pysnmp.smi import builder

SNMP_EXTEND_MIB = "NET-SNMP-EXTEND-MIB"

logger = getLogger(__name__)


class SNMPRequestTimedOut(Exception):
    pass


class SNMPInconsistentValue(Exception):
    pass


@dataclass
class SNMPResult:
    errorIndication: ErrorIndication  # True value indicates SNMP engine error
    errorStatus: str  # True value indicates SNMP PDU error
    errorIndex: int  # If nonzero, PDU error index for varBinds[errorIndex-1]
    varBinds: Tuple[ObjectType, ...]  # MIB variables returned in SNMP response


class SNMPClient:
    def __init__(self, snmp_request_timeout: float, snmp_retries: int):
        """
        :param snmp_request_timeout: The timeout for SNMP requests, in seconds
        :param snmp_retries: The number of retries for SNMP requests
        """
        self._engine = SnmpEngine()

        mib_builder = self._engine.getMibBuilder()
        mib_builder.addMibSources(builder.DirMibSource(dirname(__file__)))
        self._snmp_request_timeout = snmp_request_timeout
        self._snmp_retries = snmp_retries

    def get_system_name(self, target_ip: IPv4Address, community_string: str) -> str:
        """
        Gets the system name of the target

        :param target_ip: The IP address of the target
        :param community_string: The community string to use
        :return: A string representing the system name
        :raises SNMPRequestTimedOut: If the request times out
        :raises: Exception if the command failed
        """
        cmd = getCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(ObjectIdentity("SNMPv2-MIB", "sysName", 0)),
        )

        result = self._dispatch(cmd)
        _, name = result.varBinds[0]
        return str(name)

    def create_command(
        self, target_ip: IPv4Address, command_name: str, community_string: str, command: str
    ):
        self._try_create_command(target_ip, command_name, community_string)
        try:
            self._set_command(target_ip, command_name, community_string, command)
            self._activate_command(target_ip, command_name, community_string)
        except Exception:
            self.clear_command(target_ip, command_name, community_string)
            raise

    def _try_create_command(self, target_ip: IPv4Address, command_name: str, community_string: str):
        try:
            self._create_command(target_ip, command_name, community_string)
        except Exception:
            status = self._get_rowstatus(target_ip, command_name, community_string)
            logger.debug(f"Row status for command {command_name} is {status}")
            raise

    def _create_command(self, target_ip: IPv4Address, command_name: str, community_string: str):
        cmd = setCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(
                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name),
                "createAndWait",
            ),
            lookupNames=True,
            lookupValues=True,
        )
        self._dispatch(cmd)

    def _set_command(
        self,
        target_ip: IPv4Address,
        command_name: str,
        community_string: str,
        command: str,
    ):
        """
        Sets a command with the given name in the MIB table

        :param target_ip: The IP address of the target
        :param command_name: The name of the command to set
        :param community_string: The community string to use
        :param command: The command to set
        :raises SNMPRequestTimedOut: If the request times out
        :raises: Exception if the command failed
        """
        logger.debug(f"Using command of length: {len(command)}: {command}")
        cmd = setCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(
                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", command_name),
                "/bin/sh",
            ),
            ObjectType(
                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendArgs", command_name),
                command,
            ),
            lookupNames=True,
            lookupValues=True,
        )
        self._dispatch(cmd)

    def _activate_command(self, target_ip: IPv4Address, command_name: str, community_string: str):
        cmd = setCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(
                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name),
                "active",
            ),
        )
        self._dispatch(cmd)

    def execute_command(self, target_ip: IPv4Address, command_name: str, community_string: str):
        """
        Executes the command with the given name. The command must have been set previously

        :param target_ip: The IP address of the target
        :param command_name: The name of the command to clear
        :param community_string: The community string to use
        :raises SNMPRequestTimedOut: If the request times out
        :raises: Exception if the command failed
        """
        # Other potential NET-SNMP-EXTEND-MIB options:
        # - nsExtendStatus
        # - nsExtendOutput1Line
        # - nsExtendOutputFull
        # - nsExtendOutNumLines
        # - nsExtendOutLine
        cmd = getCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendNumEntries", 0)),
            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendCommand", command_name)),
            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendArgs", command_name)),
            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendInput", command_name)),
            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendResult", command_name)),
        )
        self._dispatch(cmd)

    def clear_command(self, target_ip: IPv4Address, command_name: str, community_string: str):
        """
        Clears the command with the given name from the MIB table. The command must have been set
        previously

        :param target_ip: The IP address of the target
        :param command_name: The name of the command to clear
        :param community_string: The community string to use
        :raises SNMPRequestTimedOut: If the request times out
        :raises: Exception if the command failed
        """
        cmd = setCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(
                ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name),
                "destroy",
            ),
        )
        self._dispatch(cmd)

    def _get_rowstatus(
        self, target_ip: IPv4Address, command_name: str, community_string: str
    ) -> str:
        cmd = getCmd(
            self._engine,
            CommunityData(community_string),
            UdpTransportTarget(
                (str(target_ip), 161),
                timeout=self._snmp_request_timeout,
                retries=self._snmp_retries,
            ),
            ContextData(),
            ObjectType(ObjectIdentity(SNMP_EXTEND_MIB, "nsExtendStatus", command_name)),
        )
        result = self._dispatch(cmd)
        vars = [
            f"{var_bind[0].prettyPrint()}={var_bind[1].prettyPrint()}"
            for var_bind in result.varBinds
        ]
        return ", ".join(vars)

    def _dispatch(self, command_iterator: Iterator[Tuple]) -> SNMPResult:
        result = SNMPResult(*next(command_iterator))

        if result.errorIndication:
            if result.errorIndication == requestTimedOut:
                raise SNMPRequestTimedOut(str(requestTimedOut))
            raise result.errorIndication
        elif result.errorStatus:
            vars = "?"
            if result.errorIndex:
                name, value = result.varBinds[int(result.errorIndex) - 1]
                vars = f"{name.prettyPrint()} = {value.prettyPrint()}"
            message = f"{result.errorStatus} at " f"{vars}"
            if str(result.errorStatus) == "inconsistentValue":
                raise SNMPInconsistentValue(message)

            raise Exception(f"Failed to execute command: {message}")

        return result
